package org.ztr.aiollama.controller;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import org.ztr.aiollama.entity.po.OllamaGenerateResponse;
import reactor.core.publisher.Mono;

/**
 * @Author: ZhaoTR
 * @Date: Created in 2025/9/20 20:59
 * @Description: 自动生成控制器
 * @Version: 1.0
 */
@RestController
@RequestMapping("/ai/generate")
public class GenerateController {
    private static final Logger logger = LogManager.getLogger(GenerateController.class);

    private final WebClient webClient = WebClient.builder()
            .baseUrl("http://localhost:11434")
            .build();

    @GetMapping("/chat")
    public Mono<String> generate(@RequestParam("prompt") String prompt) {
        return webClient.post()
                .uri("/api/generate")
                .bodyValue("""
                {
                  "model": "deepseek-r1:7b",
                  "prompt": "%s",
                  "stream": false
                }
                """.formatted(prompt))
                .retrieve()
                .bodyToMono(OllamaGenerateResponse.class)
                .map(resp -> resp.getResponse())
                .onErrorReturn("调用失败");
    }
}
